
import torch.nn as nn
import torch.nn.functional as F
from functools import partial

nonlinearity = partial(F.relu, inplace=True)


class BPNet(nn.Module):

    def __init__(self, in_seq_len=40, out_seq_len=184):
        super(BPNet, self).__init__()
        self.in_seq_len = in_seq_len
        self.out_seq_len = out_seq_len

        # fc
        self.fc1 = nn.Linear(in_seq_len, 256)
        self.relu1 = nonlinearity
        self.fc2 = nn.Linear(256, 512)
        self.relu2 = nonlinearity
        self.fc3 = nn.Linear(512, 1024)
        self.relu3 = nonlinearity
        self.fc4 = nn.Linear(1024, 2048)
        self.relu4 = nonlinearity
        self.fc5 = nn.Linear(2048, 4089)
        self.relu5 = nonlinearity
        self.fc6 = nn.Linear(4089, 2048)
        self.relu6 = nonlinearity
        self.fc7 = nn.Linear(2048, 1024)
        self.relu7 = nonlinearity
        self.fc8 = nn.Linear(1024, 512)
        self.relu8 = nonlinearity
        self.fc9 = nn.Linear(512, 256)
        self.relu9 = nonlinearity
        self.fc10 = nn.Linear(256, out_seq_len)
        self.relu10 = nonlinearity

    def forward(self, X):
        # 输入X为模型[B, 2, 20]
        batch_size = X.shape[0]

        input = X.view(batch_size, -1)

        output = self.fc1(input)
        output = self.relu1(output)

        output = self.fc2(output)
        output = self.relu2(output)

        output = self.fc3(output)
        output = self.relu3(output)

        output = self.fc4(output)
        output = self.relu4(output)

        output = self.fc5(output)
        output = self.relu5(output)

        output = self.fc6(output)
        output = self.relu6(output)

        output = self.fc7(output)
        output = self.relu7(output)

        output = self.fc8(output)
        output = self.relu8(output)

        output = self.fc9(output)
        output = self.relu9(output)

        output = self.fc10(output)

        #output = torch.tanh(output)

        # 输出处理1
        # output = self.relu10(output)
        # output = F.softmax(output, dim=1)

        # 输出处理2
        #output = torch.sigmoid(output )


        output = output.view(batch_size, 1, self.out_seq_len)

        return output